#include "MotionSegmenter.h"


//-----------------------------------------

void MotionSegmenter::openSegmentation() {

    qDebug() << "Open Segmentation";

    clusterList.clear();
    vertexToCluster.clear();

    QString filename = QFileDialog::getOpenFileName( NULL, "Open segmentation file", "", "segmentation (*.seg)" );
    QFile segfile(filename);
    if ( !segfile.open(QIODevice::ReadOnly | QIODevice::Text) ) {
        qDebug() << "File not valid";
        return;
    }
    //prepare size
    vertexToCluster.resize( _meshlist[_active]->vert.size() );

    int segcounter = 0;

    QTextStream in(&segfile);
    while(!in.atEnd()) {
        QString line = in.readLine();
        QStringList tokens = line.split(" ");
        //First two tokens contains "Segment" and number of segment: discard
        //assert(tokens.at(0) == QString("Segment"));
        //Create cluster and copy cluster in clusterlist
        std::set<int> dummyCluster;

        //then add each vertex to the cluster
        for ( int i = 2; i < tokens.size()-1; i++ ) {
            int a = tokens.at(i).toInt();
            dummyCluster.insert(tokens.at(i).toInt());
            //set vertexToCluster
            vertexToCluster[tokens.at(i).toInt()] = segcounter;
        }

        clusterList.push_back(dummyCluster);
        segcounter++;
    }

    qDebug() << "N. segmenti " << clusterList.size();

    //Coloring...
    for ( int i = 0; i < clusterList.size(); i ++ )  {

        for ( std::set<int>::iterator it = clusterList[i].begin(); it != clusterList[i].end(); ++ it ) {

            for (int j = 0; j < _meshlist.size(); j ++ )
                _meshlist[j]->vert[*it].C() = vcg::Color4b::Scatter(clusterList.size(), i);
        }
    }

    emit updateView();

}


void MotionSegmenter::openSegmentationSumner() {

    qDebug() << "Open Segmentation";

	std::vector<int> faceToCluster;

	std::vector<std::set<int> > faceClusterList;

    clusterList.clear();
	vertexToCluster.clear();
    faceToCluster.clear();

    QString filename = QFileDialog::getOpenFileName( NULL, "Open segmentation file", "", "segmentation (*.seg)" );
    QFile segfile(filename);
    if ( !segfile.open(QIODevice::ReadOnly | QIODevice::Text) ) {
        qDebug() << "File not valid";
        return;
    }
    //prepare size
    faceToCluster.resize( _meshlist[_active]->face.size() );
	vertexToCluster.resize( _meshlist[_active]->vert.size() );

    int segcounter = 0;

    QTextStream in(&segfile);
    while(!in.atEnd()) {
        QString line = in.readLine();
        QStringList tokens = line.split(" ");
        //First two tokens contains "Segment" and number of segment: discard
        //assert(tokens.at(0) == QString("Segment"));
        //Create cluster and copy cluster in clusterlist
        std::set<int> dummyCluster;

        //then add each vertex to the cluster
        for ( int i = 2; i < tokens.size()-1; i++ ) {
            int a = tokens.at(i).toInt();
            dummyCluster.insert(tokens.at(i).toInt());
            //set vertexToCluster
            faceToCluster[tokens.at(i).toInt()] = segcounter;
        }

        faceClusterList.push_back(dummyCluster);
        segcounter++;
    }

    qDebug() << "N. segmenti " << faceClusterList.size();

    //Coloring...
    for ( int i = 0; i < faceClusterList.size(); i ++ )  {

        for ( std::set<int>::iterator it = faceClusterList[i].begin(); it != faceClusterList[i].end(); ++ it ) {

            for (int j = 0; j < _meshlist.size(); j ++ )
                _meshlist[j]->face[*it].C() = vcg::Color4b::Scatter(faceClusterList.size(), i);
        }
    }

	clusterList.clear();
	clusterList.resize( faceClusterList.size() );

	for ( int i = 0; i < faceClusterList.size(); i ++ )  {

        for ( std::set<int>::iterator it = faceClusterList[i].begin(); it != faceClusterList[i].end(); ++ it ) {

            //get indices of face vertex
			for ( int j = 0; j < 3; j ++ ) {
				int vin = vcg::tri::Index( *_meshlist[_active], _meshlist[_active]->face[*it].V(j) );
				vertexToCluster[vin] = faceToCluster[*it];
				clusterList[vertexToCluster[vin]].insert(vin);
			}
        }
    }

	//Coloring...
    for ( int i = 0; i < clusterList.size(); i ++ )  {

        for ( std::set<int>::iterator it = clusterList[i].begin(); it != clusterList[i].end(); ++ it ) {

            for (int j = 0; j < _meshlist.size(); j ++ )
                _meshlist[j]->vert[*it].C() = vcg::Color4b::Scatter(clusterList.size(), i);
        }
    }

	saveSegmentation();

    emit updateView();

}

/**************************
 * Struttura del file:
 * Directory per la mesh i
 * Mesh divisa in m joints e n rigid parts:
 * -- n+m mesh, nome joint_index o rigid_index, con index progressivo
 * -- correspondence file:
 * -- la riga i-sima corrisponde al vertice i-simo e contiene coppie della serie (indice segmento, indice vertice nel segmento)
 **************************/

void MotionSegmenter::saveSegmentation() {

    qDebug() << "Saving segmentation...";

    if (_meshlist.empty()) return;

    std::vector< std::set< std::pair<int,int> > > correspondences( _meshlist[_active]->vert.size() );

    int rigid_number = clusterList.size();

    createJoints();

    int joint_number = clusterList.size() - rigid_number;

    //For each mesh...
    for ( int i = 0; i < _meshlist.size(); i ++ ) {

        //Create the directory
        QString dirname("..\\output\\pose");
        dirname.append(QString::number(i));

        //command for creating the directory
        QString command("mkdir ");
        command.append(dirname);
        qDebug() << command;
        system(command.toAscii());


        //Save the mesh
        int err = vcg::tri::io::ExporterPLY<CGMesh>::Save( *_meshlist[i], dirname.append("\\shape.ply").toAscii() );
        if(err) {
            /*printf("Error in saving %s: '%s'\n",
                   dirname.append("/shape.ply").toAscii(),
                   vcg::tri::io::Exporter<CGMesh>::ErrorMsg(err) );*/
            printf("Error in saving mesh...");
        }

        dirname.remove("\\shape.ply");

        //Save the parts, while building the correspondence map
        for ( int j = 0; j < clusterList.size(); j ++ ) {
            //For each cluster...

            //First: create a mesh, and save it (start with POINTS ONLY)
            CGMesh segment;
            vcg::tri::Allocator<CGMesh>::AddVertices(segment, clusterList[j].size());
            CGMesh::VertexIterator vit = segment.vert.begin();

            if ( j < rigid_number )
                dirname.append("\\rigid").append(QString::number(j)).append(".ply");
            else
                dirname.append("\\joint").append(QString::number(j)).append(".ply");

            //For each vertex in the cluster, create a new corresponce map
            std::map<int,int> meshToSegment;
            std::map<int,int> segmentToMesh;

            for (std::set<int>::iterator seg_it = clusterList[j].begin(); seg_it != clusterList[j].end(); ++seg_it ) {
                //Copy the coordinates of the shape vertex in the new mesh
                (*vit).P() = _meshlist[i]->vert[*seg_it].P();
                //Set correspondences between the newly created vertex and the vertex of the original mesh using indices
                meshToSegment[*seg_it] = (vit - segment.vert.begin());
                segmentToMesh[vit - segment.vert.begin()] = *seg_it;
                //Save the correspondences
                correspondences[*seg_it].insert( std::pair<int,int> (j, vit - segment.vert.begin()) );

                ++vit;  //move to next vertex
            }

            //Create the faces
            //UnMark all the elements
            vcg::tri::UnMarkAll( *_meshlist[i] );

            //for each vertex of the segment...
            for (std::set<int>::iterator seg_it = clusterList[j].begin(); seg_it != clusterList[j].end(); ++seg_it ) {
                //...we need to get the incident faces
                //get the node from the graph
                lemon::ListGraph::Node n = meshGraph.nodeFromId(*seg_it);
                //loop on the edges
                for (lemon::ListGraph::IncEdgeIt e(meshGraph, n); e != lemon::INVALID; ++e) {
                    //index of the current incident edge (need to explicit cast, otherwise index is not properly set)
                    int ei = meshGraph.id( lemon::ListGraph::Edge(e) );
                    //get the faces
                    int f0 = _edgesData[ei].f0;
                    int f1 = _edgesData[ei].f1;

                    //Face f0 not visited: visit it
                    if ( !vcg::tri::IsMarked(*_meshlist[i], &(_meshlist[i]->face[f0])) ) {
                        ///get the three vertices
                        int v0 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f0].V(0));
                        int v1 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f0].V(1));
                        int v2 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f0].V(2));
                        //if v0,v1 and v2 belong to the segment...
                        if ( meshToSegment.find(v0) != meshToSegment.end() && meshToSegment.find(v1) != meshToSegment.end() && meshToSegment.find(v2) != meshToSegment.end() ) {
                            //then add face to the segment
                            CGMesh::FaceIterator f = vcg::tri::Allocator<CGMesh>::AddFaces(segment, 1);
                            (*f).V(0) = &segment.vert[ meshToSegment[v0] ];
                            (*f).V(1) = &segment.vert[ meshToSegment[v1] ];
                            (*f).V(2) = &segment.vert[ meshToSegment[v2] ];
                        }

                        vcg::tri::Mark(*_meshlist[i], &(_meshlist[i]->face[f0]));
                    }

                    //Face f1 not visited: visit it
                    if ( !vcg::tri::IsMarked(*_meshlist[i], &(_meshlist[i]->face[f1])) ) {
                        ///get the three vertices
                        int v0 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f1].V(0));
                        int v1 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f1].V(1));
                        int v2 = vcg::tri::Index(*_meshlist[i], _meshlist[i]->face[f1].V(2));
                        //if v0,v1 and v2 belong to the same cluster...
                        //if ( vertexToCluster[v0] == vertexToCluster[v1] && vertexToCluster[v1] == vertexToCluster[v2] ) {
                        if ( meshToSegment.find(v0) != meshToSegment.end() && meshToSegment.find(v1) != meshToSegment.end() && meshToSegment.find(v2) != meshToSegment.end() ) {
                            //then add face to the segment
                            CGMesh::FaceIterator f = vcg::tri::Allocator<CGMesh>::AddFaces(segment, 1);
                            (*f).V(0) = &segment.vert[ meshToSegment[v0] ];
                            (*f).V(1) = &segment.vert[ meshToSegment[v1] ];
                            (*f).V(2) = &segment.vert[ meshToSegment[v2] ];
                        }

                        vcg::tri::Mark(*_meshlist[i], &(_meshlist[i]->face[f1]));
                    }
                }
            }

            //Save the segment
            int err = vcg::tri::io::ExporterPLY<CGMesh>::Save( segment, dirname.toAscii() );
            if(err) {
                //printf("Error in saving %s: '%s'\n", dirname.toAscii(), vcg::tri::io::Exporter<CGMesh>::ErrorMsg(err));
                printf("Error in saving...");
            }

            //garbage collector
            meshToSegment.clear();
            segmentToMesh.clear();

            //clean the directory name
            if ( j < rigid_number )
                dirname.remove(QString("\\rigid").append(QString::number(j)).append(".ply"));
            else
                dirname.remove(QString("\\joint").append(QString::number(j)).append(".ply"));

            vcg::tri::UnMarkAll( *_meshlist[i] );

        }
    }

    //Save all the correspondences on text file
    QFile file("..\\output\\correspondence.txt");
    if (!file.open(QIODevice::WriteOnly | QIODevice::Text))
        return;


    QTextStream out (&file);
    //first row:  number of meshes, number of joints, number of rigid parts, number of vertices
    out << _meshlist.size() << " " << joint_number << " " << rigid_number << " " << _meshlist[_active]->vert.size() << "\n";
    //other rows: one for each vertex
    for ( int i = 0; i < _meshlist[_active]->vert.size(); i ++ ) {
        //for each vertex write a pair (index of segment, index of vertex in that segment)
        for ( std::set< std::pair<int,int> >::iterator cit = correspondences[i].begin(); cit != correspondences[i].end(); ++ cit )
            out << (*cit).first << " " << (*cit).second << " ";
        out << "\n";
    }

    //Svuota tutto
    correspondences.clear();
}

void MotionSegmenter::createJoints() {
    qDebug() << "Searching for joints...";

    //Unmark all the elements
    vcg::tri::UnMarkAll( *_meshlist[_active] );

    std::vector< std::set<int> > joints;
    std::vector< std::set<int> > peels;

    //For each rigid part
    for ( int i = 0; i < clusterList.size(); i ++ ) {

        //visit each connected component
        for ( std::set<int>::iterator sit = clusterList[i].begin(); sit != clusterList[i].end(); ++sit ) {

            //current element already marked: skip and continue
            if ( vcg::tri::IsMarked( *_meshlist[_active], &(_meshlist[_active]->vert[*sit]) ) )
                continue;

            //not marked! let's check if it's a border edge
            bool border = false;
            lemon::ListGraph::Node n = meshGraph.nodeFromId(*sit);

            for (lemon::ListGraph::IncEdgeIt e(meshGraph, n); e != lemon::INVALID; ++e) {
                //salva nello stack quelli che appartengono a clusterA e NON sono stati ancora marcati
                lemon::ListGraph::Node u = meshGraph.u(e);
                lemon::ListGraph::Node v = meshGraph.v(e);

                if ( vertexToCluster[meshGraph.id(u)] != vertexToCluster[meshGraph.id(v)] ) {
                    border = true;
                    break;
                }
            }

            //not border... let's move on
            if (!border)  continue;

            //it's border time!
            std::set<int> dummyjoint;
            std::set<int> dummypeel;

            std::vector<int> stack;
            stack.push_back(*sit);

            while(!stack.empty()) {
                int current = stack.back();
                stack.pop_back();

                //already examined... let's move on
                if ( vcg::tri::IsMarked( *_meshlist[_active], &_meshlist[_active]->vert[current] )  )
                    continue;

                vcg::tri::Mark( *_meshlist[_active], &_meshlist[_active]->vert[current] );

                //not examined! check if it's on border
                bool border = false;
                lemon::ListGraph::Node n = meshGraph.nodeFromId(current);

                for (lemon::ListGraph::IncEdgeIt e(meshGraph, n); e != lemon::INVALID; ++e) {
                    //salva nello stack quelli che appartengono a clusterA e NON sono stati ancora marcati
                    lemon::ListGraph::Node u = meshGraph.u(e);
                    lemon::ListGraph::Node v = meshGraph.v(e);

                    if ( vertexToCluster[meshGraph.id(u)] != vertexToCluster[meshGraph.id(v)] ) {
                        border = true;
                        break;
                    }
                }

                //not on border... let's move on
                if (!border) continue;

                //it's on border! let's put the vertex and all of its neighbors in the joint cluster
                dummyjoint.insert(current);
                int radius = 3;
                std::set<int> neigh = getNeighborhood( current, radius );
                dummyjoint.insert(neigh.begin(), neigh.end());
                stack.insert(stack.end(), neigh.begin(), neigh.end());

                std::set<int> peeling = getNeighborhood( current, radius-2 );
                dummypeel.insert(peeling.begin(), peeling.end());

            }

            joints.push_back(dummyjoint);
            peels.push_back(dummypeel);
        }
    }

    for ( int i = 0; i < joints.size(); i ++ )
        clusterList.push_back( joints[i] );

    for ( int i = 0; i < peels.size(); i ++) {
        for ( std::set<int>::iterator pit = peels[i].begin(); pit != peels[i].end(); ++ pit ) {
            //_meshlist[_active]->vert[*pit].C() = vcg::Color4b::Gray;
            int c = vertexToCluster[*pit];
            clusterList[c].erase(*pit);
            vertexToCluster[*pit] = JOINT;
        }
    }

    vcg::tri::UnMarkAll( *_meshlist[_active] );
}


/**
 * Get the radius-neighborhood of central vertex
 **/
std::set<int> MotionSegmenter::getNeighborhood(int central, int radius) {

    //Distance between central and other vertices
    std::map<int, int> distance;
    //central has distance 0 from itself
    distance[central] = 0;
    //store in the stack
    std::vector<int> stack;
    stack.push_back(central);

    std::set<int> neighborhood;

    if ( radius < 0 ) return neighborhood;

    while (!stack.empty()) {
        //pop element from the stack
        int current = stack.back();
        stack.pop_back();

        //be sure that it has an element in the map
        assert( distance.find(current) != distance.end() );

        //save in the set
        neighborhood.insert(current);

        //check if we are in the feasible radius
        if ( distance[current] >= radius )
            continue;

        //we are! visit the 1-ring neighborhood of current
        lemon::ListGraph::Node n = meshGraph.nodeFromId(current);

        for (lemon::ListGraph::IncEdgeIt e(meshGraph, n); e != lemon::INVALID; ++e) {
            //one of them is different
            lemon::ListGraph::Node u = meshGraph.u(e);
            lemon::ListGraph::Node v = meshGraph.v(e);
            //find neighbor
            int neighbor = meshGraph.id(u) != current? meshGraph.id(u) : meshGraph.id(v);
            //check distance
            if ( distance.find(neighbor) == distance.end()  ) {
                distance[neighbor] = distance[current]+1;
                //put neighbor in stack
                stack.push_back(neighbor);
            }
        }
    }

    return neighborhood;
}


/**
 * Operatore di laplace beltrami
 **/
void MotionSegmenter::laplaceOperator(void) {

    laplacian = new LaplaceOperator( *_meshlist[_active] );

    laplacian->LaplaceSmoothingDirect(4, QString("cotangent"));

    emit remakeMesh();

    emit updateView();

}
